# 初始化参数

import numpy as np
import matplotlib.pyplot as plt
import sklearn              #机器学习的第三方工具包
import sklearn.datasets     #sklearn的数据集
import init_utils           #初始化

train_X, train_Y, test_X, test_Y = init_utils.load_dataset(is_plot=False)

def model(X, Y, num_iteration = 10000, learning_rate = 0.03, initialization = "random", print_cost = True, is_plot = True):

    costs = []
    layers_dims = [X.shape[0], 10, 5, 1]

    if initialization == "zeros":
        parameters = initialize_zeros(layers_dims)
    elif initialization == "random":
        parameters = initialize_random(layers_dims)
    elif initialization == "he":
        parameters = initialize_he(layers_dims)
    else:
        print("error")
        exit

    for i in range(0,num_iteration):
        a3, cache = init_utils.forward_propagation(X, parameters)
        grads = init_utils.backward_propagation(X, Y, cache)
        parameters = init_utils.update_parameters(parameters, grads, learning_rate)

        if i%1000 == 0:
            cost = init_utils.compute_loss(a3, Y)
            costs.append(cost)
            if print_cost:
                print(cost)

    if is_plot:
        plt.plot(costs)
        plt.show()

    init_utils.predict(test_X, test_Y, parameters)

    return parameters


def initialize_zeros(layers_dims):

    parameters = {}
    for i in range(1, len(layers_dims)):
        parameters["W" + str(i)] = np.zeros((layers_dims[i],layers_dims[i-1]))
        parameters["b" + str(i)] = np.zeros((layers_dims[i],1))

    for i in range(1, len(layers_dims)):
        assert(parameters["W" + str(i)].shape == (layers_dims[i],layers_dims[i-1]))
        assert(parameters["b" + str(i)].shape == (layers_dims[i],1))

    return parameters

def initialize_random(layers_dims):

    parameters = {}
    for i in range(1, len(layers_dims)):
        parameters["W" + str(i)] = np.random.randn(layers_dims[i],layers_dims[i-1])
        parameters["b" + str(i)] = np.random.randn(layers_dims[i],1)

    for i in range(1, len(layers_dims)):
        assert(parameters["W" + str(i)].shape == (layers_dims[i],layers_dims[i-1]))
        assert(parameters["b" + str(i)].shape == (layers_dims[i],1))

    return parameters

def initialize_he(layers_dims):

    parameters = {}
    for i in range(1, len(layers_dims)):
        parameters["W" + str(i)] = np.random.randn(layers_dims[i],layers_dims[i-1])*np.sqrt(2/layers_dims[i-1])
        parameters["b" + str(i)] = np.random.randn(layers_dims[i],1)*np.sqrt(2/layers_dims[i-1])

    for i in range(1, len(layers_dims)):
        assert(parameters["W" + str(i)].shape == (layers_dims[i],layers_dims[i-1]))
        assert(parameters["b" + str(i)].shape == (layers_dims[i],1))

    return parameters

parameters = model(train_X, train_Y, num_iteration = 10000, learning_rate = 0.03, initialization = "he", print_cost = True, is_plot = True)

init_utils.plot_decision_boundary(lambda x: init_utils.predict_dec(parameters, x.T), train_X, train_Y)